import os.path
import json
from preprocess.data_utils import traverse_dir
import shutil
from tqdm import tqdm


def parse_img(img_path, out_dir, image_saved_dir):
    iter_files = traverse_dir(img_path, ('jpg', 'jpeg', 'png', 'jfif', 'exif', 'gif', 'tiff', 'bmp'))
    jpg_data = {'file_num': 0, "file_suffix": ".jpg", 'file_list': []}
    has_jpg = False
    png_data = {'file_num': 0, "file_suffix": ".png", 'file_list': []}
    has_png = False
    bmp_data = {'file_num': 0, "file_suffix": ".bmp", 'file_list': []}
    has_bmp = False
    jpeg_data = {'file_num': 0, "file_suffix": ".jpeg", 'file_list': []}
    has_jpeg = False
    for fid, (file_path, file_name) in enumerate(tqdm(iter_files, desc='正在收集图片文件信息')):
        img_type = file_path.split('.')[-1]
        file_name = f'{img_type}_{fid}.{img_type}'
        shutil.copyfile(file_path, os.path.join(image_saved_dir, file_name))
        image_key = f'image_{img_type}'
        content = f'<image:{image_key}>'
        image_info = {
            image_key: {'content': '待补充...', 'path': file_name}}
        if img_type == 'jpg':
            jpg_data['file_num'] += 1
            tmp = {'file_name': file_name, 'image_info': image_info, 'table_info': {}, 'content': content}
            jpg_data['file_list'].append(tmp)
            has_jpg = True
        elif img_type == 'png':
            png_data['file_num'] += 1
            tmp = {'file_name': file_name, 'image_info': image_info, 'table_info': {}, 'content': content}
            png_data['file_list'].append(tmp)
            has_png = True
        elif img_type == 'bmp':
            bmp_data['file_num'] += 1
            tmp = {'file_name': file_name, 'image_info': image_info, 'table_info': {}, 'content': content}
            bmp_data['file_list'].append(tmp)
            has_bmp = True
        elif img_type == 'jpeg':
            jpeg_data['file_num'] += 1
            tmp = {'file_name': file_name, 'image_info': image_info, 'table_info': {}, 'content': content}
            jpeg_data['file_list'].append(tmp)
            has_jpeg = True

    if has_jpg:
        out_file = os.path.join(out_dir, 'jpg_content.json')
        with open(out_file, 'w', encoding='utf-8') as f:
            json.dump(jpg_data, f, ensure_ascii=False)
    if has_bmp:
        out_dir = os.path.join(out_dir, 'bmp_content.json')
        with open(out_dir, 'w', encoding='utf-8') as f:
            json.dump(bmp_data, f, ensure_ascii=False)

    if has_jpeg:
        out_dir = os.path.join(out_dir, 'jpeg_content.json')
        with open(out_dir, 'w', encoding='utf-8') as f:
            json.dump(jpeg_data, f, ensure_ascii=False)

    if has_png:
        out_dir = os.path.join(out_dir, 'png_content.json')
        with open(out_dir, 'w', encoding='utf-8') as f:
            json.dump(png_data, f, ensure_ascii=False)


if __name__ == '__main__':
    parse_img('../../dataset/img_files', '../out/parse_split')
